import tensorflow as tf

# 加载MNIST数据集
mnist = tf.keras.datasets.mnist
(x_train, y_train), (x_test, y_test) = mnist.load_data()

# 改变维度来适应CNN网络的输入要求
x_train4D = x_train.reshape(x_train.shape[0],28,28,1).astype('float32') 
x_test4D = x_test.reshape(x_test.shape[0],28,28,1).astype('float32')

# 数据预处理 
# *********** 请介绍数据预处理步骤 **************
x_train, x_test = x_train4D / 255.0, x_test4D / 255.0

# ***********请针对以下卷积神经网络模型的代码段给出关键注释（14至26行）**************
# 使用Sequential模型来堆叠各个层
model = tf.keras.models.Sequential([

    # 第一层卷积层
    # 32个3x3的卷积核，使用'same'填充，输入图像尺寸为28x28x1（灰度图像），激活函数为ReLU
    tf.keras.layers.Conv2D(filters=32, kernel_size=(3,3), padding='same',
                 input_shape=(28,28,1),  activation='relu'),

    # 第一层最大池化层
    # 2x2的池化窗口，步长默认为池化窗口大小
    tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),

    # 第二层卷积层
    # 64个3x3的卷积核，使用'same'填充，激活函数为ReLU
    tf.keras.layers.Conv2D(filters=64, kernel_size=(3,3), padding='same', 
    			 activation='relu'),

    # 第二层最大池化层
    # 2x2的池化窗口
    tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),

    # Dropout层，用于防止过拟合
    # 丢弃概率为0.25
    tf.keras.layers.Dropout(0.25),

    # Flatten层，将前一层的多维输出展平成一维
    tf.keras.layers.Flatten(),

    # 全连接层（密集层）
    # 128个神经元，激活函数为ReLU
    tf.keras.layers.Dense(128, activation='relu'),

    # Dropout层，用于防止过拟合
    # 丢弃概率为0.5
    tf.keras.layers.Dropout(0.5),

    # 输出层
    # 10个神经元，对应10个类别，激活函数为softmax，用于多分类
    tf.keras.layers.Dense(10,activation='softmax')
])

# 定义交叉熵损失函数
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)

# 定义优化器和学习率
optimizer = tf.keras.optimizers.Adam(learning_rate=0.001)

# 编译模型
model.compile(optimizer=optimizer,
              loss=loss_fn,
              metrics=['accuracy'])

# 训练模型
model.fit(x_train, y_train, epochs=5, validation_data=(x_test, y_test))

# 评估模型
model.evaluate(x_test, y_test)
